import sys
import cv2
import time
from PyQt5.QtWidgets import QApplication
from ultralytics import YOLO
from multiprocessing import Queue
from pathlib import Path
import json
from threading import Thread
import os


def load_config():
    """加载配置文件"""
    if getattr(sys, 'frozen', False):
        base_dir = Path(sys._MEIPASS)
    else:
        base_dir = Path(__file__).parent

    config_path = base_dir / "config.json"

    try:
        with open(config_path, "r", encoding="utf-8") as f:
            config = json.load(f)
    except FileNotFoundError:
        sys.exit(f"错误：配置文件 {config_path} 不存在")
    except json.JSONDecodeError:
        sys.exit(f"错误：配置文件格式不正确")

    model_path = Path(config["model_path"])
    if not model_path.exists():
        sys.exit(f"错误：模型文件 {model_path} 不存在")

    image_folder = config.get("image_folder", "images")
    if not Path(image_folder).exists():
        sys.exit(f"错误：图片文件夹 {image_folder} 不存在")

    return {
        "image_folder": image_folder,
        "model_path": str(model_path),
        "save_dirs": config["save_dirs"],
    }


# 保存未框选的原始帧（队列+线程）
class RawFrameSaver:
    def __init__(self, save_dir="raw_frames"):
        config = load_config()
        self.save_dir = Path(config["save_dirs"]["raw_frames"])
        self.save_dir.mkdir(exist_ok=True)
        self.frame_queue = Queue(maxsize=30)
        self.thread = Thread(target=self._save_worker, daemon=True)
        self.thread.start()

    def add_frame(self, frame, timestamp_str):
        if self.frame_queue.qsize() < 30:
            self.frame_queue.put((frame.copy(), timestamp_str))

    def _save_worker(self):
        while True:
            if self.frame_queue.empty():
                time.sleep(0.1)
                continue
            frame, timestamp_str = self.frame_queue.get()
            filename = f"{timestamp_str}.jpg"
            cv2.imwrite(str(self.save_dir / filename), frame)


# 保存带框的完整帧（直接保存，无需队列）
class FrameSaver:
    def __init__(self, save_dir="full_frames"):
        config = load_config()
        self.save_dir = Path(config["save_dirs"]["full_frames"])
        self.save_dir.mkdir(exist_ok=True)

    def save_full_frame(self, frame, timestamp_str):
        filename = f"{timestamp_str}.jpg"
        save_path = self.save_dir / filename
        cv2.imwrite(str(save_path), frame)


# 图像保存模块（保存裁剪的目标）
class ImageSaver:
    def __init__(self, save_dir="detected_objects"):
        config = load_config()
        self.save_dir = Path(config["save_dirs"]["detected_objects"])
        self.save_dir.mkdir(exist_ok=True)
        self.frame_counter = 0
        self.last_save_time = time.time()
        self.min_interval = 0.5  # 仍保留裁剪图防抖

    def save_cropped_image(self, frame, box, class_name, conf):
        current_time = time.time()
        if current_time - self.last_save_time < self.min_interval:
            return

        x1, y1, x2, y2 = map(int, box)
        cropped = frame[y1:y2, x1:x2]
        if cropped.size == 0:
            return

        timestamp = time.strftime("%Y%m%d_%H%M%S")
        filename = f"{timestamp}_{class_name}_{conf:.2f}_{self.frame_counter:04d}.jpg"
        save_path = self.save_dir / filename

        try:
            ret, buffer = cv2.imencode('.jpg', cropped, [cv2.IMWRITE_JPEG_QUALITY, 90])
            if ret:
                buffer.tofile(str(save_path))
                self.frame_counter += 1
                self.last_save_time = current_time
        except Exception as e:
            print(f"保存图像时发生错误：{str(e)}")


# 推理处理类（优化版：所有检测图都保存）
class InferenceProcessor:
    def __init__(self, model_path):
        self.model = YOLO(model_path)
        self.saver = ImageSaver()
        self.frame_saver = FrameSaver()
        self.raw_saver = RawFrameSaver()
        self.class_colors = {
            "hel": (0, 255, 0),      # 绿色
            "nohel": (0, 0, 255),    # 红色
            "wheel": (255, 0, 0),    # 蓝色
        }

    def process_frame(self, frame, timestamp_str):
        height, width, _ = frame.shape
        base_height = 720.0
        font_scale = max(0.6, height / base_height * 0.8)
        font_thickness = max(1, int(font_scale * 1.5))

        raw_frame = frame.copy()
        detected_any = False

        # ✅ 全局推理：只做一次
        results = self.model(frame, imgsz=768, conf=0.3, iou=0.45, device='cpu')

        for result in results:
            boxes = result.boxes.xyxy.cpu().numpy()
            confidences = result.boxes.conf.cpu().numpy()
            class_ids = result.boxes.cls.cpu().numpy().astype(int)

            for i, (xyxy, conf, cls_id) in enumerate(zip(boxes, confidences, class_ids)):
                if conf < 0.5:
                    continue
                class_name = result.names[cls_id]
                x1, y1, x2, y2 = map(int, xyxy)

                # --- 绘制边界框和标签（优化版：粗体、清晰）---
                color = self.class_colors.get(class_name, (255, 255, 255))  # 类别颜色

                # 绘制加粗的边界框（thickness=3）
                cv2.rectangle(frame, (x1, y1), (x2, y2), color, thickness=3)

                # 要显示的标签：类别 + 置信度
                label = f"{class_name} {conf:.2f}"

                # 字体设置：使用更粗的字体
                font = cv2.FONT_HERSHEY_DUPLEX  # 或 FONT_HERSHEY_BOLD（但 OpenCV 没有直接 BOLD，用 DUPLEX 或 SIMPLEX 加粗）
                font_scale = 0.7
                font_thickness = 2

                # 获取文本大小
                (text_width, text_height), _ = cv2.getTextSize(label, font, font_scale, font_thickness)

                # 文本背景框
                text_x, text_y = x1, max(0, y1 - 10)
                cv2.rectangle(frame,
                              (text_x, text_y - text_height - 8),
                              (text_x + text_width + 6, text_y + 4),
                              color, -1)

                # 绘制文字（白色，加粗）
                cv2.putText(frame, label,
                            (text_x + 3, text_y - 4),
                            font,
                            font_scale,
                            (255, 255, 255),  # 白色文字
                            font_thickness,
                            cv2.LINE_AA)  # 抗锯齿，更清晰

                # 保存裁剪图像
                self.saver.save_cropped_image(frame.copy(), [x1, y1, x2, y2], class_name, conf)

                detected_any = True

        # ✅ 只要检测到目标，立即保存整图和原始图（无时间限制）
        if detected_any:
            self.raw_saver.add_frame(raw_frame, "detect_" + timestamp_str)
            self.frame_saver.save_full_frame(frame.copy(), "detect_" + timestamp_str)

        return frame


# 主程序入口
if __name__ == "__main__":
    app = QApplication(sys.argv)

    config = load_config()

    # 创建保存目录
    for dir_path in config["save_dirs"].values():
        Path(dir_path).mkdir(exist_ok=True)

    # 初始化推理处理器
    processor = InferenceProcessor(
        model_path=config["model_path"]
    )

    # 读取图片文件夹
    image_folder = Path(config["image_folder"])
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}
    image_files = [f for f in image_folder.iterdir() if f.suffix.lower() in image_extensions]
    image_files.sort()

    print(f"共发现 {len(image_files)} 张图片，开始处理...")

    for img_file in image_files:
        frame = cv2.imread(str(img_file))
        if frame is None:
            print(f"无法读取图片：{img_file}")
            continue

        timestamp_str = img_file.stem + "_" + str(int(time.time() * 1000))
        print(f"正在处理：{img_file.name}")
        processed_frame = processor.process_frame(frame, timestamp_str)

    print("所有图片处理完成，程序退出。")
    sys.exit(0)